from openai import OpenAI
import os
from dotenv import load_dotenv
import time

# pip install openai dotenv


def get_api_key():
    """获取API KEY"""
    load_dotenv()
    api_key = os.getenv("DEEPSEEK_API_KEY")
    return api_key


def get_base_url(platform="deepseek"):
    """获取BASE URL"""
    if platform == "deepseek":
        # print("已经获取到了BASE_URL，平台是deepseek")
        return "https://api.deepseek.com"
    elif platform == "openai":
        # print("已经获取到了BASE_URL，平台是openai")
        return "https://api.openai.com"
    else:
        print("不支持的平台")
        return None


def ai_chat(prompt, **kwargs):
    """
    调用AI接口进行对话
    :param prompt: 提问内容
    :param kwargs: 其他参数
    :return: AI回答内容
    """

    api_key = get_api_key()
    base_url = get_base_url()
    client = OpenAI(api_key=api_key, base_url=base_url)

    prompts = f"""
    你是一个人工智能助手，你的任务是回答用户的问题。
    请用中文回答，不要用英文回答。
    用户的问题如下：{prompt}
    """

    stream = kwargs.get("stream", False)
    response = client.chat.completions.create(
        model=kwargs.get("model", "deepseek-chat"),
        messages=[
            {"role": "system", "content": "You are a helpful assistant"},
            {"role": "user", "content": f"{prompts}"},
        ],
        temperature=kwargs.get("temperature", 1.0),
        max_tokens=kwargs.get("max_tokens", 4096),
        stream=stream,
    )

    if stream:
        contents = ""
        print("\nAI回复：")
        for chunk in response:
            content = chunk.choices[0].delta.content
            print(content, end="", flush=True)
            contents += content
    else:
        contents = response.choices[0].message.content
        print("\nAI回复：")
        print(contents)

    return contents


def seconds_to_hms(seconds):
    """将秒数转换为时分秒格式"""
    # 使用divmod函数同时获取商和余数
    hours, remainder = divmod(seconds, 3600)  # 先计算小时
    minutes, seconds = divmod(remainder, 60)  # 再计算分钟和剩余秒数
    return int(hours), int(minutes), int(seconds)


if __name__ == "__main__":
    # 自定义模型参数
    custom_params = {"temperature": 0.9, "max_tokens": 150, "stream": True}

    while True:
        prompt = input("\n请输入您的问题：")

        if prompt.lower() == "exit":
            print("程序已退出")
            break

        start_time = time.time()

        # 通过**解包字典作为关键字参数
        response_content = ai_chat(prompt, **custom_params)
        print(f"AI回复：{response_content}")

        end_time = time.time()
        total_seconds = end_time - start_time
        hours, minutes, seconds = seconds_to_hms(total_seconds)
        print(f"程序运行时间：{hours}小时{minutes}分钟{seconds}秒")
